'''
Created on Oct 30, 2009

@author: mkiyer
'''

from mirna import get_mirna_target_sites
from optparse import OptionParser
from plot_utr_expr import plot_utr_expr
from veggie.sample.samplegroup import parse_samplegroups_xml, SampleGroup
import h5py
import logging
import math
import numpy as np
import operator
import os
import sys
import veggie.coverage.covdb as covdb
from veggie.sample.sampledb2 import get_sampledb
from veggie.stats.copa import calc_copa_score, copa_transform
from veggie.stats.empirical_dist import get_joint_dist, get_p_value

class UTR:
    def __init__(self):
        pass
    def __str__(self):
        return '\t'.join([self.symbol, self.accession, self.chrom, self.strand, str(self.start), str(self.end), str(self.polya)])
    
    @staticmethod        
    def read_utr_coverage_file(fhd):
        for line in fhd:
            if line is None:
                continue
            line = line.strip()
            if line is None:
                continue
            if line.startswith('#'):
                continue
            fields = line.split('\t')
            if len(fields) < 8:
                continue
            u = UTR()
            u.symbol = fields[0]
            u.accession = fields[1]
            u.chrom = fields[2]
            u.strand = fields[3]
            u.start = int(fields[4])
            u.end = int(fields[5])
            u.polya = int(fields[6])
            # for now combined exper and control
            u.proximal_covs = parse_covs_field(fields[7])
            u.distal_covs = parse_covs_field(fields[8])            
            u.proximal_covs.extend(parse_covs_field(fields[9]))
            u.distal_covs.extend(parse_covs_field(fields[10]))
            yield u         

def parse_covs_field(copa_str):
    covs = [pair.split(',') for pair in copa_str.split('|')]        
    covs = [(name, float(x)) for name, x in covs]
    return covs

def determine_outliers(ids, covs, threshold=0.75, upreg=True):
    benign_params = set(['Benign Cell Line', 'Benign Tissue'])
    cancer_params = set(['Localized Cell Line', 'Metastatic Tissue', 'Localized Tissue', 'Metastatic Cell Line'])
    test_param = 'diagnosis' 
    copa_score = calc_copa_score(covs, r=threshold)
    copa_covs = copa_transform(covs)
    outliers = []
    cancer_outliers = []
    for id, copa_cov in zip(ids, copa_covs):
        if upreg:
            if copa_cov >= copa_score:
                outliers.append(id)
                if get_sampledb().params[id][test_param] in cancer_params:
                    cancer_outliers.append(id)
        else:
            if copa_cov <= copa_score:
                outliers.append(id)
                if get_sampledb().params[id][test_param] in cancer_params:
                    cancer_outliers.append(id)            
    return copa_score, outliers, cancer_outliers

def get_joint_dist(xarr, yarr, x_binsize, y_binsize):
    '''
    given numpy arrays xarr and yarr with respective bin sizes x_binsize
    and y_binsize, compute the empirical pdf f(x,y)
    
    returns an empirical distribution matrix along with bin sizes needed
    to compute individual probabilities
    '''
    # function to generate equal sized bins over the range of an array's values
    abin = lambda a, bin_size: np.arange(0, bin_size * (1 + np.ceil(np.max(np.abs(a))/float(bin_size))), bin_size)
    mybins = (abin(xarr, x_binsize),
              abin(yarr, y_binsize))
    h, xedges, yedges = np.histogram2d(np.abs(xarr), np.abs(yarr), bins=mybins, normed=True)
    return h, x_binsize, y_binsize    
    
def get_p_value(empirical_dist, x, y):
    '''
    given an empirical distribution f(x,y) with bin sizes and 
    the values x and y, return the probability of P(X >= x, Y >= y)
    
    -the 'empirical_dist' parameter is the return value of the 
     get_joint_dist() function.  it is a 3-tuple (h, x_binsize, y_binsize).
    '''
    h, x_binsize, y_binsize = empirical_dist
    xbin = int(float(np.abs(x) / x_binsize))
    ybin = int(float(np.abs(y) / y_binsize))
    return (1.0 - (np.sum(h[0:xbin,0:ybin]) * x_binsize * y_binsize))

if __name__ == '__main__':    
    logging.basicConfig(level=logging.DEBUG)
    proximal_threshold = 0.9
    distal_threshold = 0.05
    recurrence_threshold = 1
    p_value = 0.05
    covdb_file = '/archive10/bigwig/coverage/rnaseq.hdf5'
    xml_file = '../breast_cell_lines.xml'
    exper = 'cancer'
    control = 'benign'
    raw_coverage_file = '../breast_cell_lines_raw_coverage.txt'
    do_plots = True
    
    # open coverage db
    h5file = h5py.File(covdb_file, 'r')
    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    e_sample_group = None
    c_sample_group = None
    for sgroup in parse_samplegroups_xml(xml_file):
        if sgroup.name == exper:
            e_sample_group = sgroup
        if sgroup.name == control:
            c_sample_group = sgroup
    if e_sample_group is None or c_sample_group is None:
        logging.critical("sample group not found")
        sys.exit(0)        
    e_sample_group = covdb.SampleGroupCoverage(e_sample_group, h5file)
    c_sample_group = covdb.SampleGroupCoverage(c_sample_group, h5file)


#    utrs = []
#    for u in UTR.read_utr_coverage_file(open(raw_coverage_file)):
#        # divide the tuples into separate arrays
#        proximal_ids = [x[0] for x in u.proximal_covs]
#        proximal_covs = [x[1] for x in u.proximal_covs]        
#        distal_ids = [x[0] for x in u.distal_covs]
#        distal_covs = [x[1] for x in u.distal_covs]
#        u.proximal_copa = calc_copa_score(proximal_covs, r=proximal_threshold)
#        u.distal_copa = calc_copa_score(distal_covs, r=distal_threshold)
#        utrs.append(u)
#    accepted_utrs = []
#    rejected_utrs = []
#    empirical_dist = get_joint_dist([u.proximal_copa for u in utrs],
#                                    [u.distal_copa for u in utrs], 
#                                    0.01, 0.01)
#    for u in utrs:
#        p = get_p_value(empirical_dist, u.proximal_copa, u.distal_copa)
#        if p < p_value:
#            accepted_utrs.append(u)            
#            #print u.proximal_copa, u.distal_copa, p
#        else:
#            rejected_utrs.append(u)
    # plot accepted and rejected transcripts
#    import matplotlib.pyplot as plt
#    fig = plt.figure()
#    ax = fig.add_subplot(111)
#    ax.scatter([x.proximal_copa for x in rejected_utrs], 
#               [x.distal_copa for x in rejected_utrs], s=10, c='b', marker='o', alpha=0.25)
#    if len(accepted_utrs) > 0:
#        ax.scatter([x.proximal_copa for x in accepted_utrs], 
#                   [x.distal_copa for x in accepted_utrs], s=30, c='r', marker='o', alpha=1.0)
#    plt.title('Rejected versus accepted')
#    plt.xlabel('proximal copa score')
#    plt.ylabel('distal copa score')
#    #ax.axvline(x=min_length, linewidth=1, color='r')
#    plt.show()
#    plt.close()

    debug_every = 1000
    linenum = 0
    
    proximal_copas = []
    distal_copas = []
    recurrences = []
    outlier_utrs = []
    
    for u in UTR.read_utr_coverage_file(open(raw_coverage_file)):
#        if linenum % debug_every == 0:
#            logging.debug('finished %d lines' % linenum)
        if u.strand == '+':
            mirnas = get_mirna_target_sites(u.chrom, u.polya, u.end)
        else:
            mirnas = get_mirna_target_sites(u.chrom, u.start, u.polya)
        if len(mirnas) == 0:
            continue

        # divide the tuples into separate arrays
        proximal_ids = [x[0] for x in u.proximal_covs]
        proximal_covs = [x[1] for x in u.proximal_covs]
        distal_ids = [x[0] for x in u.distal_covs]
        distal_covs = [x[1] for x in u.distal_covs]
        
        # set some RPKM limit because otherwise get tons of garbage
        if np.max(proximal_covs) < 10.0:
            continue
        if np.median(proximal_covs) < 5.0:
            continue
        if np.median(distal_covs) < 5.0:
            continue        
        
        # find overexpression outliers proximally
        proximal_copa, proximal_outliers, proximal_cancer_outliers = determine_outliers(proximal_ids, proximal_covs,
                                                                                        threshold=proximal_threshold,
                                                                                        upreg=True)
        # find underexpression outliers distally
        distal_copa, distal_outliers, distal_cancer_outliers = determine_outliers(distal_ids, distal_covs,
                                                                                  threshold=distal_threshold,
                                                                                  upreg=False)
        combined_outliers = set(proximal_outliers)
        combined_outliers.intersection_update(distal_outliers)        
        combined_cancer_outliers = set(proximal_cancer_outliers)
        combined_cancer_outliers.intersection_update(distal_cancer_outliers)
        if len(combined_outliers) == 0:
            combined_frac_cancer = 0.0
        else:
            combined_frac_cancer = float(len(combined_cancer_outliers))/len(combined_outliers)

        if len(combined_outliers) > 0:
            proximal_copas.append(proximal_copa)
            distal_copas.append(np.abs(distal_copa))
            recurrences.append(len(combined_outliers))
            u.proximal_copa = proximal_copa
            u.distal_copa = distal_copa
            outlier_utrs.append(u)

        if (len(combined_outliers) >= recurrence_threshold and
            combined_frac_cancer > 0.99):
            sys.stdout.write('%s\t%.3f\t%.3f\t%.3f\t%d\t%s\t%s\n' % 
                             (u, proximal_copa, distal_copa,
                              combined_frac_cancer, 
                              len(combined_outliers),
                              ','.join(combined_outliers), 
                              ','.join(mirnas)))
            if do_plots:
                trunc_group_cov = covdb.SampleGroupCoverage(SampleGroup('truncated', 'green', samples=combined_outliers), h5file)
                sgroups = [trunc_group_cov, c_sample_group, e_sample_group]
                name = '%s_%s' % (u.symbol, u.accession)
                title = '%s:%d-%d name=%s strand=%s' % (u.chrom, u.start, u.end, name, u.strand)            
                outfile = '%s_%s_%d-%d_%d.png' % (name, u.chrom, u.start, u.end, u.polya) 
                plot_utr_expr(u.chrom, u.start, u.end, u.strand, u.polya, sgroups,
                              outfile=outfile, 
                              title=title,
                              shade_start=u.polya)
    h5file.close()
    print len(outlier_utrs)
    
#    accepted_utrs = []
#    rejected_utrs = []
#    empirical_dist = get_joint_dist(proximal_copas, distal_copas, 0.01, 0.01)
#    for u in outlier_utrs:
#        p = get_p_value(empirical_dist, u.proximal_copa, u.distal_copa)
#        if p < p_value:
#            print u, u.proximal_copa, u.distal_copa, p
#            accepted_utrs.append(u)
#        else:
#            rejected_utrs.append(u)
    # plot outliers
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(111)    
    ax.scatter([x.proximal_copa for x in outlier_utrs], 
               [x.distal_copa for x in outlier_utrs], s=10, c='b', marker='o', alpha=0.25)
#    if len(accepted_utrs) > 0:
#        ax.scatter([x.proximal_copa for x in accepted_utrs], 
#                   [x.distal_copa for x in accepted_utrs], s=30, c='r', marker='o', alpha=1.0)
    plt.title('Outlier UTRs')
    plt.xlabel('proximal copa score')
    plt.ylabel('distal copa score')
#    #ax.axvline(x=min_length, linewidth=1, color='r')
    plt.show()
    plt.close()